{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting torch==2.2\n",
      "  Using cached torch-2.2.0-cp39-none-macosx_11_0_arm64.whl.metadata (25 kB)\n",
      "Requirement already satisfied: filelock in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (4.9.0)\n",
      "Requirement already satisfied: sympy in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (1.12)\n",
      "Requirement already satisfied: networkx in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (3.1.3)\n",
      "Requirement already satisfied: fsspec in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2) (2024.2.0)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from jinja2->torch==2.2) (2.1.5)\n",
      "Requirement already satisfied: mpmath>=0.19 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from sympy->torch==2.2) (1.3.0)\n",
      "Using cached torch-2.2.0-cp39-none-macosx_11_0_arm64.whl (59.7 MB)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.0 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0mInstalling collected packages: torch\n",
      "  Attempting uninstall: torch\n",
      "    Found existing installation: torch 1.12.1\n",
      "    Uninstalling torch-1.12.1:\n",
      "      Successfully uninstalled torch-1.12.1\n",
      "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
      "torchdata 0.4.1 requires torch==1.12.1, but you have torch 2.2.0 which is incompatible.\u001b[0m\u001b[31m\n",
      "\u001b[0mSuccessfully installed torch-2.2.0\n",
      "Collecting torchtext==0.17\n",
      "  Using cached torchtext-0.17.0-cp39-cp39-macosx_11_0_arm64.whl.metadata (7.6 kB)\n",
      "Requirement already satisfied: tqdm in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchtext==0.17) (4.66.2)\n",
      "Requirement already satisfied: requests in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchtext==0.17) (2.31.0)\n",
      "Requirement already satisfied: torch==2.2.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchtext==0.17) (2.2.0)\n",
      "Requirement already satisfied: numpy in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchtext==0.17) (1.26.4)\n",
      "Collecting torchdata==0.7.1 (from torchtext==0.17)\n",
      "  Using cached torchdata-0.7.1-cp39-cp39-macosx_11_0_arm64.whl.metadata (13 kB)\n",
      "Requirement already satisfied: filelock in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchtext==0.17) (3.13.1)\n",
      "Requirement already satisfied: typing-extensions>=4.8.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchtext==0.17) (4.9.0)\n",
      "Requirement already satisfied: sympy in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchtext==0.17) (1.12)\n",
      "Requirement already satisfied: networkx in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchtext==0.17) (2.8.5)\n",
      "Requirement already satisfied: jinja2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchtext==0.17) (3.1.3)\n",
      "Requirement already satisfied: fsspec in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torch==2.2.0->torchtext==0.17) (2024.2.0)\n",
      "Requirement already satisfied: urllib3>=1.25 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from torchdata==0.7.1->torchtext==0.17) (2.2.1)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests->torchtext==0.17) (3.3.2)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests->torchtext==0.17) (3.6)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from requests->torchtext==0.17) (2024.2.2)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from jinja2->torch==2.2.0->torchtext==0.17) (2.1.5)\n",
      "Requirement already satisfied: mpmath>=0.19 in /Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages (from sympy->torch==2.2.0->torchtext==0.17) (1.3.0)\n",
      "Using cached torchtext-0.17.0-cp39-cp39-macosx_11_0_arm64.whl (2.1 MB)\n",
      "Using cached torchdata-0.7.1-cp39-cp39-macosx_11_0_arm64.whl (4.8 MB)\n",
      "\u001b[33mDEPRECATION: pytorch-lightning 1.7.0 has a non-standard dependency specifier torch>=1.9.*. pip 24.0 will enforce this behaviour change. A possible replacement is to upgrade to a newer version of pytorch-lightning or contact the author to suggest that they release a version with a conforming dependency specifiers. Discussion can be found at https://github.com/pypa/pip/issues/12063\u001b[0m\u001b[33m\n",
      "\u001b[0mInstalling collected packages: torchdata, torchtext\n",
      "  Attempting uninstall: torchdata\n",
      "    Found existing installation: torchdata 0.4.1\n",
      "    Uninstalling torchdata-0.4.1:\n",
      "      Successfully uninstalled torchdata-0.4.1\n",
      "  Attempting uninstall: torchtext\n",
      "    Found existing installation: torchtext 0.13.1\n",
      "    Uninstalling torchtext-0.13.1:\n",
      "      Successfully uninstalled torchtext-0.13.1\n",
      "Successfully installed torchdata-0.7.1 torchtext-0.17.0\n"
     ]
    }
   ],
   "source": [
    "!pip install torch==2.2\n",
    "!pip install torchtext==0.17"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "import time\n",
    "\n",
    "import torch\n",
    "from torch import nn, Tensor\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import TransformerEncoder, TransformerEncoderLayer\n",
    "from torch.utils.data import dataset\n",
    "\n",
    "from torchtext.datasets import PennTreebank\n",
    "from torchtext.data.utils import get_tokenizer\n",
    "from torchtext.vocab import build_vocab_from_iterator\n",
    "\n",
    "torch.use_deterministic_algorithms(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Transformer(nn.Module):\n",
    "    def __init__(self, num_token, num_inputs, num_heads, num_hidden, num_layers, dropout=0.3):\n",
    "        super(Transformer, self).__init__()\n",
    "        self.model_name = 'transformer'\n",
    "        self.position_enc = PosEnc(num_inputs, dropout)\n",
    "        layers_enc = TransformerEncoderLayer(num_inputs, num_heads, num_hidden, dropout)\n",
    "        self.enc_transformer = TransformerEncoder(layers_enc, num_layers)\n",
    "        self.enc = nn.Embedding(num_token, num_inputs)\n",
    "        self.num_inputs = num_inputs\n",
    "        self.dec = nn.Linear(num_inputs, num_token)\n",
    "        self.init_params()\n",
    "\n",
    "    def init_params(self):\n",
    "        initial_rng = 0.12\n",
    "        self.enc.weight.data.uniform_(-initial_rng, initial_rng)\n",
    "        self.dec.bias.data.zero_()\n",
    "        self.dec.weight.data.uniform_(-initial_rng, initial_rng)\n",
    "\n",
    "    def forward(self, source, mask_source):\n",
    "        source = self.enc(source) * math.sqrt(self.num_inputs)\n",
    "        source = self.position_enc(source)\n",
    "        op = self.enc_transformer(source, mask_source)\n",
    "        op = self.dec(op)\n",
    "        return op\n",
    "\n",
    "def gen_sqr_nxt_mask(size):\n",
    "    msk = torch.triu(torch.ones(size, size) * float('-inf'), diagonal=1)\n",
    "    return msk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PosEnc(nn.Module):\n",
    "    def __init__(self, d_m, dropout=0.2, size_limit=5000):\n",
    "        super(PosEnc, self).__init__()\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        p_enc = torch.zeros(size_limit, 1, d_m)\n",
    "        pos = torch.arange(size_limit, dtype=torch.float).unsqueeze(1)\n",
    "        divider = torch.exp(torch.arange(0, d_m, 2).float() * (-math.log(10000.0) / d_m))\n",
    "        p_enc[:, 0, 0::2] = torch.sin(pos * divider)\n",
    "        p_enc[:, 0, 1::2] = torch.cos(pos * divider)\n",
    "        self.register_buffer('p_enc', p_enc)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        return self.dropout(x + self.p_enc[:x.size(0)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "tr_iter = PennTreebank(split='train')\n",
    "tkzer = get_tokenizer('basic_english')\n",
    "vocabulary = build_vocab_from_iterator(map(tkzer, tr_iter), specials=['<unk>'])\n",
    "vocabulary.set_default_index(vocabulary['<unk>'])\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "def process_data(raw_text):\n",
    "    numericalised_text = [torch.tensor(vocabulary(tkzer(text)), dtype=torch.long) for text in raw_text]\n",
    "    return torch.cat(tuple(filter(lambda t: t.numel() > 0, numericalised_text)))\n",
    "\n",
    "tr_iter, val_iter, te_iter = PennTreebank()\n",
    "training_text = process_data(tr_iter)\n",
    "validation_text = process_data(val_iter)\n",
    "testing_text = process_data(te_iter)\n",
    "\n",
    "def gen_batches(text_dataset, batch_size):\n",
    "    num_batches = text_dataset.size(0) // batch_size\n",
    "    text_dataset = text_dataset[:num_batches * batch_size]\n",
    "    text_dataset = text_dataset.view(batch_size, num_batches).t().contiguous()\n",
    "    return text_dataset.to(device)\n",
    "\n",
    "training_batch_size = 32\n",
    "evaluation_batch_size = 16\n",
    "\n",
    "training_data = gen_batches(training_text, training_batch_size)\n",
    "validation_data = gen_batches(validation_text, evaluation_batch_size)\n",
    "testing_data = gen_batches(testing_text, evaluation_batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_len = 64\n",
    "def return_batch(src, k):\n",
    "    sequence_length = min(max_seq_len, len(src) - 1 - k)\n",
    "    sequence_data = src[k:k+sequence_length]\n",
    "    sequence_label = src[k+1:k+1+sequence_length].reshape(-1)\n",
    "    return sequence_data, sequence_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/Ashish.Jha/anaconda3/envs/python39/lib/python3.9/site-packages/torch/nn/modules/transformer.py:286: UserWarning: enable_nested_tensor is True, but self.use_nested_tensor is False because encoder_layer.self_attn.batch_first was not True(use batch_first for better inference performance)\n",
      "  warnings.warn(f\"enable_nested_tensor is True, but self.use_nested_tensor is False because {why_not_sparsity_fast_path}\")\n"
     ]
    }
   ],
   "source": [
    "num_tokens = len(vocabulary) # vocabulary size\n",
    "embedding_size = 256 # dimension of embedding layer\n",
    "num_hidden_params = 256 # transformer encoder's hidden (feed forward) layer dimension\n",
    "num_layers = 2 # num of transformer encoder layers within transformer encoder\n",
    "num_heads = 2 # num of heads in (multi head) attention models\n",
    "dropout = 0.25 # value (fraction) of dropout\n",
    "loss_func = nn.CrossEntropyLoss()\n",
    "lrate = 4.0 # learning rate\n",
    "transformer_model = Transformer(num_tokens, embedding_size, num_heads, num_hidden_params, num_layers, \n",
    "                                     dropout).to(device)\n",
    "optim_module = torch.optim.SGD(transformer_model.parameters(), lr=lrate)\n",
    "sched_module = torch.optim.lr_scheduler.StepLR(optim_module, 1.0, gamma=0.88)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model():\n",
    "    transformer_model.train()\n",
    "    loss_total = 0.\n",
    "    time_start = time.time()\n",
    "    mask_source = gen_sqr_nxt_mask(max_seq_len).to(device)\n",
    "    num_batches = len(training_data) // max_seq_len\n",
    "    for b, i in enumerate(range(0, training_data.size(0) - 1, max_seq_len)):\n",
    "        train_data_batch, train_label_batch = return_batch(training_data, i)\n",
    "        sequence_length = train_data_batch.size(0)\n",
    "        if sequence_length != max_seq_len:  # only on last batch\n",
    "            mask_source = mask_source[:sequence_length, :sequence_length]\n",
    "        op = transformer_model(train_data_batch, mask_source)\n",
    "        loss_curr = loss_func(op.view(-1, num_tokens), train_label_batch)\n",
    "        optim_module.zero_grad()\n",
    "        loss_curr.backward()\n",
    "        torch.nn.utils.clip_grad_norm_(transformer_model.parameters(), 0.6)\n",
    "        optim_module.step()\n",
    "\n",
    "        loss_total += loss_curr.item()\n",
    "        interval = 100\n",
    "        if b % interval == 0 and b > 0:\n",
    "            loss_interval = loss_total / interval\n",
    "            time_delta = time.time() - time_start\n",
    "            print(f\"epoch {ep}, {b}/{len(training_data)//max_seq_len} batches, training loss {loss_interval:.2f}, training perplexity {math.exp(loss_interval):.2f}\")\n",
    "            loss_total = 0\n",
    "            time_start = time.time()\n",
    "\n",
    "def eval_model(eval_model_obj, eval_data_source):\n",
    "    eval_model_obj.eval() \n",
    "    loss_total = 0.\n",
    "    mask_source = gen_sqr_nxt_mask(max_seq_len).to(device)\n",
    "    with torch.no_grad():\n",
    "        for j in range(0, eval_data_source.size(0) - 1, max_seq_len):\n",
    "            eval_data, eval_label = return_batch(eval_data_source, j)\n",
    "            sequence_length = eval_data.size(0)\n",
    "            if sequence_length != max_seq_len:\n",
    "                mask_source = mask_source[:sequence_length, :sequence_length]\n",
    "            op = eval_model_obj(eval_data, mask_source)\n",
    "            op_flat = op.view(-1, num_tokens)\n",
    "            loss_total += sequence_length * loss_func(op_flat, eval_label).item()\n",
    "    return loss_total / (len(eval_data_source) - 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch 1, 100/451 batches, training loss 7.86, training perplexity 2588.27\n",
      "epoch 1, 200/451 batches, training loss 6.44, training perplexity 628.79\n",
      "epoch 1, 300/451 batches, training loss 6.03, training perplexity 416.68\n",
      "epoch 1, 400/451 batches, training loss 5.83, training perplexity 340.21\n",
      "\n",
      "epoch 1, validation loss 5.65, validation perplexity 282.91\n",
      "\n",
      "epoch 2, 100/451 batches, training loss 5.68, training perplexity 291.70\n",
      "epoch 2, 200/451 batches, training loss 5.58, training perplexity 265.24\n",
      "epoch 2, 300/451 batches, training loss 5.45, training perplexity 232.82\n",
      "epoch 2, 400/451 batches, training loss 5.37, training perplexity 215.16\n",
      "\n",
      "epoch 2, validation loss 5.40, validation perplexity 220.48\n",
      "\n",
      "epoch 3, 100/451 batches, training loss 5.37, training perplexity 213.94\n",
      "epoch 3, 200/451 batches, training loss 5.31, training perplexity 203.13\n",
      "epoch 3, 300/451 batches, training loss 5.20, training perplexity 181.45\n",
      "epoch 3, 400/451 batches, training loss 5.14, training perplexity 170.07\n",
      "\n",
      "epoch 3, validation loss 5.27, validation perplexity 193.97\n",
      "\n",
      "epoch 4, 100/451 batches, training loss 5.18, training perplexity 177.30\n",
      "epoch 4, 200/451 batches, training loss 5.14, training perplexity 170.42\n",
      "epoch 4, 300/451 batches, training loss 5.03, training perplexity 153.42\n",
      "epoch 4, 400/451 batches, training loss 4.97, training perplexity 144.43\n",
      "\n",
      "epoch 4, validation loss 5.16, validation perplexity 174.62\n",
      "\n",
      "epoch 5, 100/451 batches, training loss 5.04, training perplexity 154.86\n",
      "epoch 5, 200/451 batches, training loss 5.01, training perplexity 150.22\n",
      "epoch 5, 300/451 batches, training loss 4.91, training perplexity 135.80\n",
      "epoch 5, 400/451 batches, training loss 4.85, training perplexity 127.64\n",
      "\n",
      "epoch 5, validation loss 5.10, validation perplexity 164.12\n",
      "\n"
     ]
    }
   ],
   "source": [
    "min_validation_loss = float(\"inf\")\n",
    "eps = 5\n",
    "best_model_so_far = None\n",
    "\n",
    "for ep in range(1, eps + 1):\n",
    "    ep_time_start = time.time()\n",
    "    train_model()\n",
    "    validation_loss = eval_model(transformer_model, validation_data)\n",
    "    print()\n",
    "    print(f\"epoch {ep:}, validation loss {validation_loss:.2f}, validation perplexity {math.exp(validation_loss):.2f}\")\n",
    "    print()\n",
    "\n",
    "    if validation_loss < min_validation_loss:\n",
    "        min_validation_loss = validation_loss\n",
    "        best_model_so_far = transformer_model\n",
    "\n",
    "    sched_module.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "testing loss 5.03, testing perplexity 152.75\n"
     ]
    }
   ],
   "source": [
    "testing_loss = eval_model(best_model_so_far, testing_data)\n",
    "print(f\"testing loss {testing_loss:.2f}, testing perplexity {math.exp(testing_loss):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
